# -*- coding: UTF-8 -*-
"""
    @Author:YTQ
    @Time: 2022/7/20 19:45
    Description:使用模型进行图片分类
    
"""

import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dataset import DetectDataSet
from CNNet import MyCNN


def detect(model_path, image_path, data_transform, dogAct, catAct, batch_size, workers):
    '''
        model_path: 模型地址
        image_path：分类图片文件夹
        data_transform：数据集转换设置
        dogAct：够对应数字
        catAct：猫对应数字
        batch_size：训练所抓取的数据样本数量
        workers：加载数据线程数
    '''
    # load images
    dataset = DetectDataSet(image_path, data_transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
                            num_workers=workers)
    # load model
    cnn = MyCNN()
    # cuda
    if torch.cuda.is_available() is True:
        print('Cuda is available!')
        cnn = cnn.cuda()
    cnn.load_state_dict(torch.load(model_path))

    for img, label in dataloader:  # 循环读取封装后的数据集，其实就是调用了数据集中的__getitem__()方法，只是返回数据格式进行了一次封装
        img = Variable(img)
        out = cnn(img)[0]  # 计算网络输出值，就是输入网络一个图像数据，输出猫和狗的概率，调用了网络中的forward()方法
        if out[0, 0] > out[0, 1]:  # 猫的概率大于狗
            return dogAct
        else:  # 猫的概率小于狗
            return catAct


